from openai import OpenAI
from openai.types.chat.chat_completion import Choice


def doubao_chat(api_key: str, text: str) -> Choice:
    # 初始化Ark客户端
    client = OpenAI(
        base_url="https://ark.cn-beijing.volces.com/api/v3",
        api_key=api_key,
    )

    response = client.chat.completions.create(
        model="doubao-seed-1-6-flash-250615",
        messages=[
            {
                "role": "user",
                "content": [
                    {"type": "text", "text": text},
                ],
            }
        ],
    )

    return response.choices[0]


if __name__ == "__main__":
    import fire
    fire.Fire(doubao_chat)
